#!/usr/bin/env Python
# coding = utf-8

'''
zpasswd.py r 读取密码
zpasswd.py w 写入密码
'''
import pickle
import sys
import argparse
import string
import random

#.eg: python zpasswd.py -m w/r web qq or python zpasswd.py -m g -l 10
parser = argparse.ArgumentParser()
parser.add_argument('-m', action='store', dest='method', help='操作方式 w：写密码，r：查找密码， g：生成密码')
parser.add_argument('-c', action='store', dest='category', help='类别')
parser.add_argument('-i', action='store', dest='title', help='查找的名称')
parser.add_argument('-l', type=int, action='store', dest='length', help='密码长度')

args = parser.parse_args()

#读文件
def read_passwd():
	try:
		with open('zpasswd.pass', 'rb') as f:
			return pickle.load(f)
	except Exception as e:
		return False
		raise e

#写文件
def write_passwd(user_data):
	try:
		with open('zpasswd.pass', 'wb') as f:
			pickle.dump(user_data, f)
			return True
	except Exception as e:
		return False
		raise e


'''
数据结构：
{category:
	{title:{'user':user,
			'passwd':passwd
			}
	}
}
'''
#密码类，直接继承字典
class PassWord(dict):
	def __init__(self, category, title, user, passwd):
		self.category = category
		self.title = title
		self[title] = {'user': user, 'passwd': passwd} 

	#静态方法
	@staticmethod
	def get_passwd():
		user_data = read_passwd() if read_passwd() else {}
		if user_data != {}:
			return user_data.get(args.category).get(args.title)
		else:
			return None	


	def set_passwd(self):
		user_data = read_passwd() if read_passwd() else {}
		#如果类别层字典不存在则创建一个空字典。
		if user_data.get(self.category) == None:
			user_data[self.category] = {}
		user_data[self.category][self.title] = self[self.title]
		#写入结果返回
		if write_passwd(user_data):
			print('ok!')
			return True
		else:
			print('not ok!')
			return False

	#类方法
	@classmethod
	def get_item(cls):
		cls.item = ['title', 'user', 'passwd']
		return cls.item


class WebPassWord(PassWord):
	#继承时候，初始化函数要写齐全父类的参数+子类的参数
	def __init__(self, category, title, user, passwd, url):
		#初始化父类的参数
		super().__init__(category, title, user, passwd)
		self[title]['url'] = url


	@classmethod
	def get_item(cls):
		super().get_item()
		cls.item.append('url')
		return cls.item


class ServerPassWord(PassWord):
	
	def __init__(self, category, title, user, passwd, ip, hostname):
		super().__init__(category, title, user, passwd)
		self[title]['ip'] = ip
		self[title]['hostname'] = hostname


	@classmethod
	def get_item(cls):
		super().get_item()
		cls.item.append('ip')
		cls.item.append('hostname')
		return cls.item



#生成随机密码
def gen_passwd(length):
	chars = string.ascii_letters + string.digits + '!@#$%^&*()'
	return ''.join([random.choice(chars) for i in range(length)])

		
def handle_passwd():
		passwd_args = {}
		user_choice = input('1 web, 2 server, 0 退出')
		pwd_class = {
		'1':('web', WebPassWord),
		'2':('server', ServerPassWord)
		}
		#可以将类名赋值给变量，使用变量调用类
		category = pwd_class[user_choice][0]
		handle_class = pwd_class[user_choice][1]
		for i in handle_class.get_item():
			passwd_args[i] = input(i + ':')
		#实例化类，使用**表达式使用字典自动赋值
		a = handle_class(category=category, **passwd_args)
		a.set_passwd()


def main():
	if args.method == 'r':
		print(PassWord.get_passwd())
	elif args.method == 'w':
		handle_passwd()
	elif args.method == 'g':
		print(gen_passwd(args.length))
	else:
		print('退出')
		sys.exit(0)


if __name__ == '__main__':
	main()